import os
import shutil
import time
from typing import Dict
from chatad.roster import Roster
from chatad.utils import log_visualize

class ChatEnvConfig:
    def __init__(self, clear_structure,
                 background_prompt):
        self.clear_structure = clear_structure  # 是否清除WareHouse中路径中的缓存文件
        self.background_prompt = background_prompt  # 每次对话时需要添加的固定prompt，在ChatChainConfig内


    def __str__(self):
        string = ""
        string += "ChatEnvConfig.clear_structure: {}\n".format(self.clear_structure)
        string += "ChatEnvConfig.background_prompt: {}\n".format(self.background_prompt)
        return string


class ChatEnv:
    def __init__(self, chat_env_config: ChatEnvConfig):
        self.config = chat_env_config
        self.roster: Roster = Roster()  # 获取多智能体列表，并可判断对话的多智能体是否存在
        self.env_dict = {
            "directory": "",
            "task_prompt": "",
            "task_description": "",
            "modality": "",
            "advertising_suggestions":"",
            "ideas": "",
            "crawled_data":"",
            "data_analysis": "",
            "market_analysis":"",
            "placement_recommendations":"",#广告投放结论
            "ppt_content":"",
            "raw_data":"",    # market analysis这个阶段中使用
            "total_data_rows":"", #market analysis的数据行数
            "available_files": [],       # 可用数据文件列表
            "classified_files": [],      # 任务相关的数据文件
            "classified_data": {},       # 任务相关的数据内容
        }



    def set_directory(self, directory):  # directory：项目文件夹名
        assert len(self.env_dict['directory']) == 0
        self.env_dict['directory'] = directory
        if os.path.exists(self.env_dict['directory']) and len(os.listdir(directory)) > 0: # 源目录已存在
            new_directory = "{}.{}".format(directory, time.strftime("%Y%m%d%H%M%S", time.localtime()))
            shutil.copytree(directory, new_directory)  # 重新创建一个目录，将源目录里的内容移动到新目录
            print("{} Copied to {}".format(directory, new_directory))
        if os.path.exists(self.env_dict['directory']):  # 如果源目录已存在，则删除这个目录，再重新创建
            shutil.rmtree(self.env_dict['directory'])
            os.mkdir(self.env_dict['directory'])
            print("{} Created".format(directory))
        else:
            os.mkdir(self.env_dict['directory'])

    def recruit(self, agent_name: str): # 将使用的多智能体存入列表
        self.roster._recruit(agent_name)

    def exist_employee(self, agent_name: str) -> bool:
        return self.roster._exist_employee(agent_name)

    def write_meta(self) -> None:
        directory = self.env_dict['directory']

        if not os.path.exists(directory):
            os.mkdir(directory)
            print("{} Created.".format(directory))

        meta_filename = "meta.txt"
        with open(os.path.join(directory, meta_filename), "w", encoding="utf-8") as writer:
            writer.write("{}:\n{}\n\n".format("Task", self.env_dict['task_prompt']))
            writer.write("{}:\n{}\n\n".format("Config", self.config.__str__()))
            writer.write("{}:\n{}\n\n".format("Roster", ", ".join(self.roster.agents)))
            writer.write("{}:\n{}\n\n".format("Modality", self.env_dict['modality']))
            writer.write("{}:\n{}\n\n".format("Ideas", self.env_dict['ideas']))

       